import torch 
import torch.nn as nn 

class Conv2d(nn.Module):
    def __init__(self, nin, nout, ks, stride=1):
        super().__init__() 
        pad = (ks-1)//2
        self.layers = nn.Sequential(
            nn.Conv2d(nin, nout, ks, stride, padding=pad), 
            nn.BatchNorm2d(nout), 
            nn.ReLU(), 
        )
    def forward(self, x):
        x = self.layers(x) 
        return x 
class Model(nn.Module):
    def __init__(self):
        super().__init__() 
        self.layers1 = nn.Sequential(
            Conv2d(3, 16, 3, 2), 
            Conv2d(16, 32, 3, 2), 
            Conv2d(32, 64, 3, 2), 
        )      
        self.layers2 = nn.Sequential(
            Conv2d(64, 64, 3, 2), 
        )
        self.layers3 = nn.Sequential(
            Conv2d(64, 64, 3, 2), 
        )
        self.out1 = nn.Conv2d(64, 4+1+80, 1) 
        self.out2 = nn.Conv2d(64, 4+1+80, 1) 
        self.out3 = nn.Conv2d(64, 4+1+80, 1) 
    def forward(self, x):
        y1 = self.layers1(x)
        y2 = self.layers2(y1) 
        y3 = self.layers3(y2)
        
        return self.out1(y1), self.out2(y2), self.out3(y3)  
model = Model() 
model.train()

image = torch.randn([1, 3, 416, 416])
y1, y2, y3 = model(image) 
print(y1.shape, y2.shape, y3.shape)

# [416, 416] 
# - [52, 52]  /8 # 10, 10 -> 
# - [26, 26]  /16 
# - [13, 13]  /32  感受野依次增大
